import graphviz

dot = graphviz.Digraph(comment='Two-Layer Unidirectional RNN', engine='dot')
dot.attr(rankdir='TD')  # 设置为从上到下布局 (Top-Down)

# 定义节点样式
node_style = 'filled, rounded'
special_node_style = 'filled, rounded, style="filled", color=lightblue'

# 1. Inputs
with dot.subgraph(name='cluster_inputs') as c:
    c.attr(style='filled', color='lightgrey', label='Inputs')
    c.node('x', 'x (batch, seq_len, input_size)', shape='box', style=node_style)
    c.node('h0', 'h0 (num_layers, batch, hidden_size)', shape='box', style=node_style)

# 2. RNN Layer 1
with dot.subgraph(name='cluster_layer1') as c:
    c.attr(style='filled', color='lightyellow', label='RNN Layer 1')
    c.node('h0_1', 'h0[0]', shape='ellipse', style=special_node_style)
    c.node('rnn1_1', 'RNN Cell 1\n(Time Step 1)', shape='box', style=node_style)
    c.node('rnn1_t', 'RNN Cell 1\n(Time Step t)', shape='box', style=node_style)
    c.node('rnn1_T', 'RNN Cell 1\n(Time Step T)', shape='box', style=node_style)
    c.node('h1_1', 'h1_1', shape='ellipse', style=special_node_style)
    c.node('h1_t', 'h1_t', shape='ellipse', style=special_node_style)
    c.node('h1_T', 'h1_T', shape='ellipse', style=special_node_style)

    c.edge('x', 'rnn1_1', label='x[:, 0, :]')  # 第一个时间步的输入
    c.edge('h0_1', 'rnn1_1')
    c.edge('rnn1_1', 'h1_1')
    c.edge('h1_1', 'rnn1_t', label='Feedback')  # 隐藏状态反馈
    c.edge('x', 'rnn1_t', label='x[:, t, :]')
    c.edge('rnn1_t', 'h1_t')
    c.edge('h1_t', 'rnn1_T',label='Feedback')
    c.edge('x', 'rnn1_T', label='x[:, T-1, :]')
    c.edge('rnn1_T', 'h1_T')
    c.node('out1','out1',shape='box', style=node_style)
    c.edge('h1_1', 'out1', label='h1_1, h1_2, ..., h1_T', style='dashed') #不可见
    c.edge('h1_t', 'out1', style='dashed')#不可见
    c.edge('h1_T', 'out1', style='dashed')#不可见


# 3. RNN Layer 2
with dot.subgraph(name='cluster_layer2') as c:
    c.attr(style='filled', color='lightblue', label='RNN Layer 2')
    c.node('h0_2', 'h0[1]', shape='ellipse', style=special_node_style)
    c.node('rnn2_1', 'RNN Cell 2\n(Time Step 1)', shape='box', style=node_style)
    c.node('rnn2_t', 'RNN Cell 2\n(Time Step t)', shape='box', style=node_style)
    c.node('rnn2_T', 'RNN Cell 2\n(Time Step T)', shape='box', style=node_style)
    c.node('h2_1', 'h2_1', shape='ellipse', style=special_node_style)
    c.node('h2_t', 'h2_t', shape='ellipse', style=special_node_style)
    c.node('h2_T', 'h2_T', shape='ellipse', style=special_node_style)

    c.edge('out1', 'rnn2_1', label='h1_1')  # 第一层的输出作为第二层的输入
    c.edge('h0_2', 'rnn2_1')
    c.edge('rnn2_1', 'h2_1')
    c.edge('h2_1', 'rnn2_t', label='Feedback')
    c.edge('out1', 'rnn2_t', label='h1_t')
    c.edge('rnn2_t', 'h2_t')
    c.edge('h2_t', 'rnn2_T',label='Feedback')
    c.edge('out1', 'rnn2_T', label='h1_T')
    c.edge('rnn2_T', 'h2_T')

# 4. Outputs
with dot.subgraph(name='cluster_outputs') as c:
    c.attr(style='filled', color='lightgreen', label='Outputs')
    c.node('out', 'out (batch, seq_len, hidden_size)', shape='box', style=node_style)
    c.node('hn', 'hn (num_layers, batch, hidden_size)', shape='box', style=node_style)
    c.node('out_last', 'out[:, -1, :]\n(batch, hidden_size)', shape='box', style=special_node_style)
    c.node('hn_last', 'hn[-1]\n(batch, hidden_size)', shape='box', style=special_node_style)

# 连接
dot.edge('h2_1', 'out')
dot.edge('h2_t', 'out')
dot.edge('h2_T', 'out')

dot.edge('h1_T', 'hn', label='hn[0]')
dot.edge('h2_T', 'hn', label='hn[1]')

dot.edge('out', 'out_last', style='dashed')
dot.edge('hn', 'hn_last', style='dashed')

# 注释
dot.node('note1', '''
out: 包含第二层 (最后一层) RNN 在每个时间步的输出
hn: 包含每一层 RNN 在最后一个时间步的隐藏状态
''', shape='note', fontsize='10')

#dot.render('rnn_flowchart', view=True, format='png')
dot.render('rnn_flowchart', format='png')

print(dot.source)
